package xyz.thoughtset.viewer.executor.blocks.executor;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.lang.NonNull;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import xyz.thoughtset.viewer.common.ai.model.entity.ModelParam;
import xyz.thoughtset.viewer.common.ai.model.factory.ModelFactory;
import xyz.thoughtset.viewer.common.ai.model.service.ModelParamService;
import xyz.thoughtset.viewer.common.exc.exceptions.ExecException;
import xyz.thoughtset.viewer.modules.step.entity.AISupportBody;
import xyz.thoughtset.viewer.executor.blocks.executor.AbstractBlockExecutor;
import xyz.thoughtset.viewer.modules.step.entity.BlockTypeEnum;
import xyz.thoughtset.viewer.modules.step.entity.block.BlockInfo;
import xyz.thoughtset.viewer.modules.step.entity.block.EleParam;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Slf4j
@SuppressWarnings({"unchecked","SpringJavaInjectionPointsAutowiringInspection"})
public abstract class AbstractAISupportBlockExecutor<T extends AISupportBody> extends AbstractBlockExecutor<T> {
    @Autowired
    protected ModelParamService modelParamService;
    @Autowired
    protected ModelFactory modelFactory;


    protected static String renderPrompt(String template, Map params){
        boolean paramsNotEmpty = !ObjectUtils.isEmpty(params);
        if ( !paramsNotEmpty || !StringUtils.hasText(template)){
            return template;
        }
        PromptTemplate promptTemplate = PromptTemplate.builder()
                    .template(template)
                    .variables(params)
                    .build();
        return promptTemplate.render();
    }

    protected HashMap filterDataAsMapForPrompt( T body, Map<String, Object> params, ExpressionParser parser, StandardEvaluationContext context){
        HashMap<String, Object> resultMaps = new HashMap<>(params);
        List<EleParam> eleParams = body.getDataParams();
        if (Objects.nonNull(eleParams)){
            for (EleParam blockParam : eleParams){
                Object value = "";
                if(StringUtils.hasLength(blockParam.getDataExp())){
                    String exp =blockParam.getDataExp();
                    exp = blockParam.getDataExp().startsWith("#")?exp:"#"+exp;
                    value = parser.parseExpression(exp).getValue(context);
                }
                String key = blockParam.getParamId();
                if (StringUtils.hasText(key) && !ObjectUtils.isEmpty(value)){
                    if (!(value instanceof String)){
                        try {
                            value = objectMapper.writeValueAsString(value);
                        } catch (Exception e) {
                            log.error("objectMapper error", value);
                            e.printStackTrace();
                            value = "";
                        }
                    }
                    resultMaps.put(key, value);
                }
            }
        }
        return resultMaps;
    }

    protected void saveMegToChatMemory(String chatId, List<Message> messages, ChatMemory chatMemory) {
        chatMemory.add(chatId, messages);

        List<Message> memoryMessage = chatMemory.get(chatId);

        long assistantCount = memoryMessage.stream()
                .filter(m -> m.getMessageType() == MessageType.ASSISTANT)
                .count();
        long toolCount = memoryMessage.stream()
                .filter(m -> m.getMessageType() == MessageType.TOOL)
                .count();

        if (assistantCount != toolCount) {
            log.error("消息不成对{}，重置",memoryMessage);
            chatMemory.clear(chatId);
            chatMemory.add(chatId, memoryMessage.subList(1, memoryMessage.size()));
        }
    }
    @SneakyThrows
    protected ModelParam loadModelParam(@NonNull AISupportBody aiBody) throws ExecException {
        String modelId = aiBody.getModelId();
        if (!StringUtils.hasText(modelId)){
            throw new Exception("模型参数未配置");
        }
        ModelParam modelParam = modelParamService.selectDetail(modelId);
        if (modelParam == null){
            throw new Exception("模型参数不存在");
        }
        return modelParam;
    }

    protected static String fillPrompt(String template, ExpressionParser parser,StandardEvaluationContext context) {
        if (!StringUtils.hasText(template)){
            return template;
        }
        try {
            return parser.parseExpression(template).getValue(context, String.class);
        }catch (Exception e){
            return template;
        }
    }

    protected static String fillAndRenderPromptToStr(String template, ExpressionParser parser,StandardEvaluationContext context,Map resultMaps) {
        return renderPrompt(
                fillPrompt(template, parser, context),resultMaps
        );
    }

}
